#include "conn_tab.h"
#include <linux/ktime.h>

#define BAD_CELL {BAD_STATE, SERVER, KEEP_TIME} 									// a bad state of tcp 
static connection_t* connTablePtr = NULL; 											// the connection table 
static cell_t tcpStateMachine[NUM_OF_STATES][NUM_OF_COMMANDS] = { { BAD_CELL } };   // tcp machine  state

// Visual representation of the tcp state machine 

//static cell_t tcpStateMachine[NUM_OF_STATES][NUM_OF_COMMANDS] = {
			         	    /*NONE	                 	| SYN                       | ACK                      |SYN_ACK                          |FIN                                |FIN_SYN                        |FIN_ACK                     |FIN_SYN_ACK                   |RESET                      */
/////*TCP_CONN_SYN_SENT*/   { TCP_CONN_SYN_SENT,BOTH },	BAD_CELL,                  BAD_CELL,                 {TCP_CONN_SYN_ACK,SERVER},       {TCP_CONN_CLOSING,BOTH},          {TCP_CONN_CLOSING,BOTH},        {TCP_CONN_CLOSING,BOTH},   {TCP_CONN_CLOSING,SERVER},	 {TCP_CONN_SYN_SENT,BOTH},			
/////*TCP_CONN_SYN_ACK*/    {TCP_CONN_ESTAB,BOTH},    	{TCP_CONN_ESTAB,BOTH},      {TCP_CONN_ESTAB,BOTH},     {TCP_CONN_ESTAB,BOTH},           {TCP_CONN_CLOSING,BOTH},          {TCP_CONN_CLOSING,BOTH},        {TCP_CONN_CLOSING,BOTH},   {TCP_CONN_CLOSING,BOTH},       {TCP_CONN_SYN_ACK,BOTH},
/////*TCP_CONN_ESTAB*/      {TCP_CONN_ESTAB,BOTH},    	{TCP_CONN_ESTAB,BOTH},      {TCP_CONN_ESTAB,BOTH},     {TCP_CONN_ESTAB,BOTH},           {TCP_CONN_CLOSING,BOTH},          {TCP_CONN_CLOSING,BOTH},        {TCP_CONN_CLOSING,BOTH},   {TCP_CONN_CLOSING,BOTH},      {TCP_CONN_ESTAB,BOTH},
/////*TCP_CONN_CLOSSING*/	 BAD_CELL,                 BAD_CELL,                  {TCP_CONN_CLOSING,BOTH},  {TCP_CONN_CLOSING,BOTH},        {TCP_CONN_CLOSING,BOTH},          {TCP_CONN_CLOSING,BOTH},        {TCP_CONN_CLOSING,BOTH},   {TCP_CONN_CLOSING,BOTH},      {TCP_CONN_CLOSING,BOTH},

//};	



void initStateMachine(void){	
	//Row 0
	tcpStateMachine[TCP_CONN_SYN_SENT-1][SYN_ACK].nextState = TCP_CONN_SYN_ACK;
	tcpStateMachine[TCP_CONN_SYN_SENT-1][SYN_ACK].side = SERVER;	
	tcpStateMachine[TCP_CONN_SYN_SENT-1][SYN_ACK].seconds = FIVE_SECONDS;	
	
	tcpStateMachine[TCP_CONN_SYN_SENT-1][FIN].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_SYN_SENT-1][FIN].side = BOTH;		
	tcpStateMachine[TCP_CONN_SYN_SENT-1][FIN].seconds = FIVE_SECONDS;	
	
	tcpStateMachine[TCP_CONN_SYN_SENT-1][FIN_ACK].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_SYN_SENT-1][FIN_ACK].side = BOTH;	
	tcpStateMachine[TCP_CONN_SYN_SENT-1][FIN_ACK].seconds = FIVE_SECONDS;
		
	tcpStateMachine[TCP_CONN_SYN_SENT-1][FIN_SYN].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_SYN_SENT-1][FIN_SYN].side = BOTH;		
	tcpStateMachine[TCP_CONN_SYN_SENT-1][FIN_SYN].seconds = FIVE_SECONDS;
	
	tcpStateMachine[TCP_CONN_SYN_SENT-1][FIN_SYN_ACK].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_SYN_SENT-1][FIN_SYN_ACK].side = SERVER;	
	tcpStateMachine[TCP_CONN_SYN_SENT-1][FIN_SYN_ACK].seconds = FIVE_SECONDS;
	
	tcpStateMachine[TCP_CONN_SYN_SENT-1][RESET].nextState = TCP_CONN_SYN_SENT;
	tcpStateMachine[TCP_CONN_SYN_SENT-1][RESET].side = BOTH;
	tcpStateMachine[TCP_CONN_SYN_SENT-1][RESET].seconds = RESET_TIME;
	
	//Row 1
	tcpStateMachine[TCP_CONN_SYN_ACK-1][NONE].nextState = TCP_CONN_ESTAB;
	tcpStateMachine[TCP_CONN_SYN_ACK-1][NONE].side = BOTH;		
	tcpStateMachine[TCP_CONN_SYN_ACK-1][NONE].seconds = FIVE_MINUTES;
	
	tcpStateMachine[TCP_CONN_SYN_ACK-1][SYN].nextState = TCP_CONN_ESTAB;
	tcpStateMachine[TCP_CONN_SYN_ACK-1][SYN].side = BOTH;
	tcpStateMachine[TCP_CONN_SYN_ACK-1][SYN].seconds = FIVE_MINUTES;
	
	tcpStateMachine[TCP_CONN_SYN_ACK-1][ACK].nextState = TCP_CONN_ESTAB;
	tcpStateMachine[TCP_CONN_SYN_ACK-1][ACK].side = BOTH;
	tcpStateMachine[TCP_CONN_SYN_ACK-1][ACK].seconds = FIVE_MINUTES;
	
	tcpStateMachine[TCP_CONN_SYN_ACK-1][SYN_ACK].nextState = TCP_CONN_ESTAB;
	tcpStateMachine[TCP_CONN_SYN_ACK-1][SYN_ACK].side = BOTH;
	tcpStateMachine[TCP_CONN_SYN_ACK-1][SYN_ACK].seconds = FIVE_MINUTES;
	
	tcpStateMachine[TCP_CONN_SYN_ACK-1][FIN].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_SYN_ACK-1][FIN].side = BOTH;		
	tcpStateMachine[TCP_CONN_SYN_ACK-1][FIN].seconds = FIVE_SECONDS;
	
	tcpStateMachine[TCP_CONN_SYN_ACK-1][FIN_ACK].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_SYN_ACK-1][FIN_ACK].side = BOTH;	
	tcpStateMachine[TCP_CONN_SYN_ACK-1][FIN_ACK].seconds = FIVE_SECONDS;
	
	tcpStateMachine[TCP_CONN_SYN_ACK-1][FIN_SYN].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_SYN_ACK-1][FIN_SYN].side = BOTH;		
	tcpStateMachine[TCP_CONN_SYN_ACK-1][FIN_SYN].seconds = FIVE_SECONDS;
	
	tcpStateMachine[TCP_CONN_SYN_ACK-1][FIN_SYN_ACK].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_SYN_ACK-1][FIN_SYN_ACK].side = BOTH;	
	tcpStateMachine[TCP_CONN_SYN_ACK-1][FIN_SYN_ACK].seconds = FIVE_SECONDS;
	
	tcpStateMachine[TCP_CONN_SYN_ACK-1][RESET].nextState = TCP_CONN_SYN_ACK;
	tcpStateMachine[TCP_CONN_SYN_ACK-1][RESET].side = BOTH;
	tcpStateMachine[TCP_CONN_SYN_ACK-1][RESET].seconds = RESET_TIME;
	
	//Row 2
	tcpStateMachine[TCP_CONN_ESTAB-1][NONE].nextState = TCP_CONN_ESTAB;
	tcpStateMachine[TCP_CONN_ESTAB-1][NONE].side = BOTH;		
	tcpStateMachine[TCP_CONN_ESTAB-1][NONE].seconds = FIVE_MINUTES;		
	
	tcpStateMachine[TCP_CONN_ESTAB-1][SYN].nextState = TCP_CONN_ESTAB;
	tcpStateMachine[TCP_CONN_ESTAB-1][SYN].side = BOTH;
	tcpStateMachine[TCP_CONN_ESTAB-1][SYN].seconds = FIVE_MINUTES;
	
	tcpStateMachine[TCP_CONN_ESTAB-1][ACK].nextState = TCP_CONN_ESTAB;
	tcpStateMachine[TCP_CONN_ESTAB-1][ACK].side = BOTH;
	tcpStateMachine[TCP_CONN_ESTAB-1][ACK].seconds = FIVE_MINUTES;
	
	tcpStateMachine[TCP_CONN_ESTAB-1][SYN_ACK].nextState = TCP_CONN_ESTAB;
	tcpStateMachine[TCP_CONN_ESTAB-1][SYN_ACK].side = BOTH;
	tcpStateMachine[TCP_CONN_ESTAB-1][SYN_ACK].seconds = FIVE_MINUTES;
	
	tcpStateMachine[TCP_CONN_ESTAB-1][FIN].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_ESTAB-1][FIN].side = BOTH;		
	tcpStateMachine[TCP_CONN_ESTAB-1][FIN].seconds = FIVE_SECONDS;
	
	tcpStateMachine[TCP_CONN_ESTAB-1][FIN_ACK].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_ESTAB-1][FIN_ACK].side = BOTH;	
	tcpStateMachine[TCP_CONN_ESTAB-1][FIN_ACK].seconds = FIVE_SECONDS;
	
	tcpStateMachine[TCP_CONN_ESTAB-1][FIN_SYN].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_ESTAB-1][FIN_SYN].side = BOTH;		
	tcpStateMachine[TCP_CONN_ESTAB-1][FIN_SYN].seconds = FIVE_SECONDS;
	
	tcpStateMachine[TCP_CONN_ESTAB-1][FIN_SYN_ACK].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_ESTAB-1][FIN_SYN_ACK].side = BOTH;	
	tcpStateMachine[TCP_CONN_ESTAB-1][FIN_SYN_ACK].seconds = FIVE_SECONDS;
	
	tcpStateMachine[TCP_CONN_ESTAB-1][RESET].nextState = TCP_CONN_ESTAB;
	tcpStateMachine[TCP_CONN_ESTAB-1][RESET].side = BOTH;
	tcpStateMachine[TCP_CONN_ESTAB-1][RESET].seconds = RESET_TIME;
	//Row 3
	tcpStateMachine[TCP_CONN_CLOSING-1][ACK].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_CLOSING-1][ACK].side = BOTH;
	tcpStateMachine[TCP_CONN_CLOSING-1][ACK].seconds = KEEP_TIME;
	
	tcpStateMachine[TCP_CONN_CLOSING-1][SYN_ACK].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_CLOSING-1][SYN_ACK].side = BOTH;
	tcpStateMachine[TCP_CONN_CLOSING-1][SYN_ACK].seconds = KEEP_TIME;
	
	tcpStateMachine[TCP_CONN_CLOSING-1][FIN].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_CLOSING-1][FIN].side = BOTH;		
	tcpStateMachine[TCP_CONN_CLOSING-1][FIN].seconds = KEEP_TIME;
	
	tcpStateMachine[TCP_CONN_CLOSING-1][FIN_ACK].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_CLOSING-1][FIN_ACK].side = BOTH;	
	tcpStateMachine[TCP_CONN_CLOSING-1][FIN_ACK].seconds = KEEP_TIME;
	
	tcpStateMachine[TCP_CONN_CLOSING-1][FIN_SYN].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_CLOSING-1][FIN_SYN].side = BOTH;		
	tcpStateMachine[TCP_CONN_CLOSING-1][FIN_SYN].seconds = KEEP_TIME;
	
	tcpStateMachine[TCP_CONN_CLOSING-1][FIN_SYN_ACK].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_CLOSING-1][FIN_SYN_ACK].side = BOTH;	
	tcpStateMachine[TCP_CONN_CLOSING-1][FIN_SYN_ACK].seconds = KEEP_TIME;
	
	tcpStateMachine[TCP_CONN_CLOSING-1][RESET].nextState = TCP_CONN_CLOSING;
	tcpStateMachine[TCP_CONN_CLOSING-1][RESET].side = BOTH;	
	tcpStateMachine[TCP_CONN_CLOSING-1][RESET].seconds = RESET_TIME;
}

connection_t* findConversation(rawPacket_t* packet){
	int i;
	unsigned long now = get_seconds();
	for(i=0;i<CONNECTION_TABLE_ENTRIES;++i){
		if(connTablePtr[i].expires < now){
			continue;
		}
		if(connTablePtr[i].cli_ip == packet->source_ip && connTablePtr[i].ser_ip == packet->dest_ip &&
			connTablePtr[i].cli_port == packet->source_port && connTablePtr[i].ser_port == packet->dest_port){
			if(connTablePtr[i].expires < now) continue;			
			return &connTablePtr[i];
		}
		if(connTablePtr[i].cli_ip == packet->dest_ip && connTablePtr[i].ser_ip == packet->source_ip &&
			connTablePtr[i].cli_port == packet->dest_port && connTablePtr[i].ser_port == packet->source_port){
			if(connTablePtr[i].expires < now) continue;			
			return &connTablePtr[i];
		}
	}
	return CONV_NOT_FOUND;
}

connection_t* addConversation(rawPacket_t* packet){
	int i;
	unsigned long now = get_seconds();
	for(i=0;i<CONNECTION_TABLE_ENTRIES;++i){
		if(connTablePtr[i].expires < now){
			//found an empty slot
			memset(&connTablePtr[i],0,sizeof(connection_t));
			connTablePtr[i].cli_ip = packet->source_ip;
			connTablePtr[i].ser_ip = packet->dest_ip;
			connTablePtr[i].cli_port = packet->source_port;
			connTablePtr[i].ser_port = packet->dest_port;
			connTablePtr[i].expires = get_seconds() + FIVE_SECONDS;
			connTablePtr[i].state = TCP_CONN_SYN_SENT;			
			return &connTablePtr[i];
		}
	}
	return CONN_TABLE_FULL;
}

tcp_conn_t nextState(tcp_conn_t currentState, tcp_commands_t command, side_t side){	
	if(currentState < TCP_CONN_SYN_SENT  || currentState > TCP_CONN_CLOSING){ 		
		return BAD_STATE; // state out of bounds
	}	
	currentState--;	//current state starts at 1
	if(command < NONE || command >= NUM_OF_COMMANDS){ 		
		return BAD_STATE; // command out of bounds
	}	
	if(tcpStateMachine[currentState][command].side == side || tcpStateMachine[currentState][command].side == BOTH){		
		return tcpStateMachine[currentState][command].nextState;
	}	
	return BAD_STATE;		
}

tcp_commands_t getCommand(tcpHeaderParams_t params){
	tcp_commands_t tempCommand = NONE;
	if(params.rst)
		return RESET;
	if(params.syn)
		tempCommand += SYN;
	if(params.ack)
		tempCommand += ACK;
	if(params.fin)
		tempCommand += FIN;	
	if(tempCommand >= NUM_OF_COMMANDS)	
		return BAD_COMMAND;	
	return tempCommand;
}

int numOfConversations(void){
	int counter = 0;
	int i;
	unsigned long now = get_seconds();
	for(i=0;i<CONNECTION_TABLE_ENTRIES;i++){
		if(connTablePtr[i].expires >= now)
			counter++;
	} 
	return counter++;
}

void setState(connection_t* conv, tcp_conn_t state, tcp_commands_t command){
	tcp_conn_t prevState = conv->state;	
	if(tcpStateMachine[prevState-1][command].seconds != KEEP_TIME){
		if(tcpStateMachine[prevState-1][command].seconds != RESET_TIME){
			conv->expires = get_seconds() + tcpStateMachine[prevState-1][command].seconds;
		}else{
			conv->expires = tcpStateMachine[prevState-1][command].seconds;
		}
	}
	conv->state = state; // update state	
}

void clearConnectionTable(void){	
	memset(connTablePtr,0,sizeof(connection_t)*CONNECTION_TABLE_ENTRIES);
}

void printConversationTable(void){
	int i;
	printk("---------- CONNECTION TABLE --------------\n");
	for(i=0;i<CONNECTION_TABLE_ENTRIES;++i){
		if(get_seconds() > connTablePtr[i].expires) continue;
		printk("%d: cli_ip %d ser_ip %d state %d\n",i,connTablePtr[i].cli_ip,connTablePtr[i].ser_ip,connTablePtr[i].state);
	}
}
